# 导入pandas库，用于数据处理
import pandas as pd
# 导入numpy库，用于数值计算
import numpy as np
# 导入matplotlib库，用于绘图
import matplotlib
# 导入警告模块，用于忽略警告信息
import warnings
# 导入绘图模块，用于可视化
import matplotlib.pyplot as plt
import io
import base64
from fastapi.responses import HTMLResponse
from typing import List, Optional, Dict, Any
import os
import platform
import sys

# 检测并设置合适的中文字体
def set_chinese_font():
    # 检测操作系统
    system = platform.system()
    
    # 根据不同操作系统设置不同的中文字体
    if system == 'Windows':
        plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'SimSun', 'Arial Unicode MS']
    elif system == 'Darwin':  # macOS
        plt.rcParams['font.sans-serif'] = ['PingFang SC', 'STHeiti', 'Heiti TC', 'Arial Unicode MS']
    else:  # Linux或其他系统
        plt.rcParams['font.sans-serif'] = ['WenQuanYi Micro Hei', 'DejaVu Sans', 'Arial Unicode MS']
    
    # 兜底方案，添加一个通用的备选字体列表
    font_list = plt.rcParams['font.sans-serif']
    font_list.extend(['Noto Sans CJK SC', 'Noto Sans CJK TC', 'Noto Sans CJK JP', 'WenQuanYi Zen Hei'])
    plt.rcParams['font.sans-serif'] = font_list
    
    # 正常显示负号
    plt.rcParams['axes.unicode_minus'] = False
    return True

# 调用函数设置中文字体
set_chinese_font()

# 更新字体大小
matplotlib.rcParams.update({'font.size': 12})
# 设置绘图风格为默认白色背景
plt.style.use('default')  # 使用默认的白色背景样式
# 设置网格不可见
plt.rcParams['axes.grid'] = False
# 忽略警告信息
warnings.filterwarnings('ignore')

class DataFrameHandler:
    """
    处理DataFrame数据的类
    实现导入文件、数据操作和可视化展示
    """
    
    @staticmethod
    def load_file(file_path: str, file_type: str = None) -> pd.DataFrame:
        """
        导入数据文件，支持csv和excel格式
        
        参数:
            file_path: 文件路径
            file_type: 文件类型，如果为None则自动从扩展名判断
        
        返回:
            加载的DataFrame
        """
        if file_type is None:
            # 从文件扩展名判断类型
            file_ext = os.path.splitext(file_path)[1].lower()
            
            if file_ext in ['.csv', '.txt']:
                return pd.read_csv(file_path, encoding='utf-8')
            elif file_ext in ['.xls', '.xlsx']:
                return pd.read_excel(file_path)
            else:
                raise ValueError(f"不支持的文件类型: {file_ext}")
        else:
            if file_type.lower() == 'csv':
                return pd.read_csv(file_path, encoding='utf-8')
            elif file_type.lower() in ['excel', 'xls', 'xlsx']:
                return pd.read_excel(file_path)
            else:
                raise ValueError(f"不支持的文件类型: {file_type}")

    @staticmethod
    def get_df_info(df: pd.DataFrame) -> Dict[str, Any]:
        """
        获取DataFrame的基本信息
        
        参数:
            df: 要分析的DataFrame
            
        返回:
            包含DataFrame信息的字典
        """
        info = {
            "shape": df.shape,
            "columns": df.columns.tolist(),
            "dtypes": {col: str(dtype) for col, dtype in zip(df.columns, df.dtypes)},
            "has_null": df.isnull().any().to_dict(),
            "describe": df.describe().to_dict()
        }
        return info
    
    @staticmethod
    def filter_data(df: pd.DataFrame, conditions: Dict[str, Any]) -> pd.DataFrame:
        """
        根据条件筛选数据
        
        参数:
            df: 要筛选的DataFrame
            conditions: 筛选条件字典，格式为 {"column": {"operator": "value"}}
            
        返回:
            筛选后的DataFrame
        """
        filtered_df = df.copy()
        
        for col, condition in conditions.items():
            if col not in df.columns:
                continue
                
            for op, value in condition.items():
                if op == "eq":
                    filtered_df = filtered_df[filtered_df[col] == value]
                elif op == "gt":
                    filtered_df = filtered_df[filtered_df[col] > value]
                elif op == "lt":
                    filtered_df = filtered_df[filtered_df[col] < value]
                elif op == "ge":
                    filtered_df = filtered_df[filtered_df[col] >= value]
                elif op == "le":
                    filtered_df = filtered_df[filtered_df[col] <= value]
                elif op == "contains":
                    filtered_df = filtered_df[filtered_df[col].astype(str).str.contains(value)]
        
        return filtered_df
    
    @staticmethod
    def group_data(df: pd.DataFrame, group_by: List[str], agg_dict: Dict[str, str]) -> pd.DataFrame:
        """
        对数据进行分组聚合
        
        参数:
            df: 要分组的DataFrame
            group_by: 分组列名列表
            agg_dict: 聚合方式字典，格式为 {"column": "agg_method"}
            
        返回:
            分组聚合后的DataFrame
        """
        return df.groupby(group_by).agg(agg_dict).reset_index()
    
    @staticmethod
    def plot_bar(df: pd.DataFrame, x: str, y: str, title: str = "柱状图", y_label: str = "数量") -> str:
        """
        绘制柱状图
        
        参数:
            df: 数据DataFrame
            x: x轴列名
            y: y轴列名
            title: 图表标题
            y_label: y轴标签，默认为"数量"
            
        返回:
            Base64编码的图像字符串
        """
        # 确保中文字体设置
        set_chinese_font()
        plt.figure(figsize=(10, 6))
        
        # 检查数据类型并进行适当处理
        if pd.api.types.is_numeric_dtype(df[y]):
            # 如果有太多唯一值的分类变量，则进行分组
            if pd.api.types.is_object_dtype(df[x]) and len(df[x].unique()) > 20:
                # 取前20个值进行展示
                top_values = df.groupby(x)[y].sum().nlargest(20).index
                filtered_df = df[df[x].isin(top_values)]
                ax = filtered_df.plot(kind='bar', x=x, y=y)
            else:
                ax = df.plot(kind='bar', x=x, y=y)
        else:
            # 如果y不是数值型，尝试计数分析
            value_counts = df[x].value_counts()
            value_counts.plot(kind='bar')
            
        plt.title(title)
        # 设置x轴和y轴标签
        plt.xlabel(x)
        plt.ylabel(y_label)  # 使用传入的y_label参数作为y轴标签
        plt.tight_layout()
        
        # 将图表转换为Base64编码
        buffer = io.BytesIO()
        plt.savefig(buffer, format='png')
        buffer.seek(0)
        img_str = base64.b64encode(buffer.read()).decode('utf-8')
        plt.close()
        
        return img_str
    
    @staticmethod
    def plot_line(df: pd.DataFrame, x: str, y: List[str], title: str = "折线图") -> str:
        """
        绘制折线图
        
        参数:
            df: 数据DataFrame
            x: x轴列名
            y: y轴列名列表
            title: 图表标题
            
        返回:
            Base64编码的图像字符串
        """
        # 确保中文字体设置
        set_chinese_font()
        plt.figure(figsize=(10, 6))
        
        # 检查数据类型
        if pd.api.types.is_datetime64_any_dtype(df[x]) or pd.api.types.is_numeric_dtype(df[x]):
            # 日期或数值型可以直接绘制
            ax = df.plot(kind='line', x=x, y=y)
        else:
            # 对于分类变量，可能需要特殊处理
            temp_df = df.copy()
            # 转换为分类型并指定顺序，确保正确排序
            temp_df[x] = pd.Categorical(temp_df[x], categories=temp_df[x].unique(), ordered=True)
            temp_df = temp_df.sort_values(x)
            
            for col in y:
                if pd.api.types.is_numeric_dtype(temp_df[col]):
                    plt.plot(temp_df[x], temp_df[col], marker='o', label=col)
                else:
                    # 如果y也不是数值型，可以尝试转换或跳过
                    try:
                        temp_df[col] = pd.to_numeric(temp_df[col], errors='coerce')
                        plt.plot(temp_df[x], temp_df[col], marker='o', label=col)
                    except:
                        continue
            
            plt.legend()
            plt.xticks(rotation=45)
            
        plt.title(title)
        # 设置x轴和y轴标签
        plt.xlabel(x)
        if len(y) == 1:
            plt.ylabel(y[0])
        else:
            plt.ylabel('数值')
        plt.tight_layout()
        
        # 将图表转换为Base64编码
        buffer = io.BytesIO()
        plt.savefig(buffer, format='png')
        buffer.seek(0)
        img_str = base64.b64encode(buffer.read()).decode('utf-8')
        plt.close()
        
        return img_str
    
    @staticmethod
    def plot_pie(df: pd.DataFrame, values: str, labels: str, title: str = "饼图") -> str:
        """
        绘制饼图
        
        参数:
            df: 数据DataFrame
            values: 值列名
            labels: 标签列名
            title: 图表标题
            
        返回:
            Base64编码的图像字符串
        """
        # 确保中文字体设置
        set_chinese_font()
        plt.figure(figsize=(10, 8))
        
        # 检查数据类型
        if not pd.api.types.is_numeric_dtype(df[values]):
            # 如果values不是数值型，尝试转换
            try:
                df[values] = pd.to_numeric(df[values], errors='coerce')
            except:
                # 如果无法转换，可以使用计数
                value_counts = df[labels].value_counts()
                plt.pie(value_counts, labels=value_counts.index, autopct='%1.1f%%', shadow=False, startangle=90)
                plt.axis('equal')
                plt.title(title)
                
                # 将图表转换为Base64编码
                buffer = io.BytesIO()
                plt.savefig(buffer, format='png')
                buffer.seek(0)
                img_str = base64.b64encode(buffer.read()).decode('utf-8')
                plt.close()
                
                return img_str
        
        # 如果标签值太多，进行分组
        if len(df[labels].unique()) > 10:
            # 取前10个值，其他归为"其他"类别
            top_values = df.groupby(labels)[values].sum().nlargest(9).index
            
            # 创建分组数据
            grouped_data = pd.DataFrame()
            grouped_data['label'] = ['其他'] + list(top_values)
            
            # 计算值
            others_sum = df[~df[labels].isin(top_values)][values].sum()
            top_sums = [df[df[labels] == val][values].sum() for val in top_values]
            grouped_data['value'] = [others_sum] + top_sums
            
            # 绘制饼图
            plt.pie(grouped_data['value'], labels=grouped_data['label'], autopct='%1.1f%%', shadow=False, startangle=90)
        else:
            # 分组计算总和
            grouped_df = df.groupby(labels)[values].sum().reset_index()
            plt.pie(grouped_df[values], labels=grouped_df[labels], autopct='%1.1f%%', shadow=False, startangle=90)
        
        plt.axis('equal')
        plt.title(title)
        
        # 将图表转换为Base64编码
        buffer = io.BytesIO()
        plt.savefig(buffer, format='png')
        buffer.seek(0)
        img_str = base64.b64encode(buffer.read()).decode('utf-8')
        plt.close()
        
        return img_str
    
    @staticmethod
    def plot_scatter(df: pd.DataFrame, x: str, y: str, title: str = "散点图") -> str:
        """
        绘制散点图
        
        参数:
            df: 数据DataFrame
            x: x轴列名
            y: y轴列名
            title: 图表标题
            
        返回:
            Base64编码的图像字符串
        """
        # 确保中文字体设置
        set_chinese_font()
        plt.figure(figsize=(10, 6))
        
        # 检查x和y的数据类型
        x_numeric = pd.api.types.is_numeric_dtype(df[x])
        y_numeric = pd.api.types.is_numeric_dtype(df[y])
        
        if not x_numeric or not y_numeric:
            # 对于非数值型数据，尝试转换
            temp_df = df.copy()
            
            if not x_numeric:
                try:
                    temp_df[x] = pd.to_numeric(temp_df[x], errors='coerce')
                except:
                    # 如果无法转换为数值，使用分类编码
                    categories = pd.Categorical(temp_df[x]).codes
                    plt.scatter(categories, temp_df[y] if y_numeric else pd.Categorical(temp_df[y]).codes)
                    
                    # 设置x轴刻度标签为原始分类值
                    unique_vals = temp_df[x].unique()
                    plt.xticks(range(len(unique_vals)), unique_vals, rotation=45)
                    
                    if not y_numeric:
                        # 如果y也不是数值，设置y轴刻度
                        unique_vals_y = temp_df[y].unique()
                        plt.yticks(range(len(unique_vals_y)), unique_vals_y)
                    
                    plt.title(title)
                    plt.xlabel(x)
                    plt.ylabel(y)
                    plt.grid(True)
                    plt.tight_layout()
                    
                    # 将图表转换为Base64编码
                    buffer = io.BytesIO()
                    plt.savefig(buffer, format='png')
                    buffer.seek(0)
                    img_str = base64.b64encode(buffer.read()).decode('utf-8')
                    plt.close()
                    
                    return img_str
            
            if not y_numeric:
                try:
                    temp_df[y] = pd.to_numeric(temp_df[y], errors='coerce')
                except:
                    # 这种情况在前面的代码中已经处理了
                    pass
            
            plt.scatter(temp_df[x], temp_df[y], alpha=0.6)
        else:
            plt.scatter(df[x], df[y], alpha=0.6)
        
        plt.title(title)
        plt.xlabel(x)
        plt.ylabel(y)
        plt.grid(True)
        plt.tight_layout()
        
        # 将图表转换为Base64编码
        buffer = io.BytesIO()
        plt.savefig(buffer, format='png')
        buffer.seek(0)
        img_str = base64.b64encode(buffer.read()).decode('utf-8')
        plt.close()
        
        return img_str
    
    @staticmethod
    def plot_to_html(img_base64: str) -> str:
        """
        将Base64编码的图像转换为HTML图像标签
        
        参数:
            img_base64: Base64编码的图像字符串
            
        返回:
            HTML图像标签
        """
        return f'<img src="data:image/png;base64,{img_base64}" alt="Data Visualization" />'
    
    @staticmethod
    def correlation_heatmap(df: pd.DataFrame, title: str = "相关性热力图") -> str:
        """
        绘制相关性热力图
        
        参数:
            df: 数据DataFrame
            title: 图表标题
            
        返回:
            Base64编码的图像字符串
        """
        # 计算相关性矩阵，只包括数值型列
        numeric_df = df.select_dtypes(include=['number'])
        corr_matrix = numeric_df.corr()
        
        # 确保中文字体设置
        set_chinese_font()
        plt.figure(figsize=(10, 8))
        plt.imshow(corr_matrix, cmap='coolwarm', interpolation='none', aspect='auto')
        plt.colorbar(label='相关系数')
        plt.title(title)
        
        # 添加坐标标签
        plt.xticks(range(len(corr_matrix.columns)), corr_matrix.columns, rotation=45)
        plt.yticks(range(len(corr_matrix.columns)), corr_matrix.columns)
        
        # 添加相关系数值
        for i in range(len(corr_matrix.columns)):
            for j in range(len(corr_matrix.columns)):
                plt.text(j, i, f'{corr_matrix.iloc[i, j]:.2f}', 
                         ha='center', va='center', color='black')
        
        # 添加轴标签
        plt.xlabel('特征')
        plt.ylabel('特征')
        plt.tight_layout()
        
        # 将图表转换为Base64编码
        buffer = io.BytesIO()
        plt.savefig(buffer, format='png')
        buffer.seek(0)
        img_str = base64.b64encode(buffer.read()).decode('utf-8')
        plt.close()
        
        return img_str
    
    @staticmethod
    def time_series_analysis(df: pd.DataFrame, date_col: str, value_col: str, 
                          freq: str = 'M', title: str = "时间序列分析") -> str:
        """
        时间序列分析与可视化
        
        参数:
            df: 数据DataFrame
            date_col: 日期列名
            value_col: 值列名
            freq: 重采样频率，如'D'表示按天，'M'表示按月
            title: 图表标题
            
        返回:
            Base64编码的图像字符串
        """
        # 确保日期列为datetime类型
        df = df.copy()
        df[date_col] = pd.to_datetime(df[date_col])
        
        # 设置日期列为索引
        df.set_index(date_col, inplace=True)
        
        # 按指定频率重采样并计算均值
        resampled = df[value_col].resample(freq).mean()
        
        # 确保中文字体设置
        set_chinese_font()
        plt.figure(figsize=(12, 6))
        resampled.plot()
        plt.title(f"{title} - {value_col}")
        # 设置x轴和y轴标签
        plt.xlabel(date_col)
        plt.ylabel(value_col)
        plt.grid(True)
        plt.tight_layout()
        
        # 将图表转换为Base64编码
        buffer = io.BytesIO()
        plt.savefig(buffer, format='png')
        buffer.seek(0)
        img_str = base64.b64encode(buffer.read()).decode('utf-8')
        plt.close()
        
        return img_str

# 提供一个简单的测试函数
def test_visualization():
    """
    测试可视化功能
    """
    # 创建一个简单的DataFrame
    df = pd.DataFrame({
        'category': ['A', 'B', 'C', 'D', 'E'],
        'values': [23, 48, 15, 32, 29]
    })
    
    handler = DataFrameHandler()
    img_base64 = handler.plot_bar(df, 'category', 'values', "测试柱状图")
    html_img = handler.plot_to_html(img_base64)
    
    return html_img

# 示例：加载CSV文件并生成可视化
def analyze_csv_example(file_path: str) -> Dict[str, str]:
    """
    分析CSV文件并生成多种可视化
    
    参数:
        file_path: CSV文件路径
        
    返回:
        包含不同可视化HTML的字典
    """
    handler = DataFrameHandler()
    try:
        df = handler.load_file(file_path)
        
        results = {
            "info": str(handler.get_df_info(df)),
            "visualizations": {}
        }
        
        # 仅对数值列进行可视化
        numeric_cols = df.select_dtypes(include=['number']).columns.tolist()
        
        if len(numeric_cols) >= 2:
            # 散点图（前两个数值列）
            scatter_img = handler.plot_scatter(df, numeric_cols[0], numeric_cols[1], f"{numeric_cols[0]} vs {numeric_cols[1]}")
            results["visualizations"]["scatter"] = handler.plot_to_html(scatter_img)
            
            # 相关性热力图
            corr_img = handler.correlation_heatmap(df, "相关性分析")
            results["visualizations"]["correlation"] = handler.plot_to_html(corr_img)
        
        # 如果有分类列和数值列，生成柱状图
        if len(df.select_dtypes(include=['object']).columns) > 0 and len(numeric_cols) > 0:
            cat_col = df.select_dtypes(include=['object']).columns[0]
            num_col = numeric_cols[0]
            
            # 对分类列进行分组统计
            grouped_df = df.groupby(cat_col)[num_col].mean().reset_index()
            bar_img = handler.plot_bar(grouped_df, cat_col, num_col, f"{cat_col}按{num_col}统计")
            results["visualizations"]["bar"] = handler.plot_to_html(bar_img)
        
        return results
    
    except Exception as e:
        return {"error": str(e)}

if __name__ == "__main__":
    # 测试可视化
    html_result = test_visualization()
    print("生成的HTML图像标签:", html_result[:100] + "...")
    
    # 测试CSV分析
    results = analyze_csv_example("./vip_info_100.csv")
    if "error" not in results:
        print("成功分析CSV文件并生成可视化")
    else:
        print("分析CSV文件时出错:", results["error"]) 